import cv2 
import os 
import numpy as np
from driver.hsv_find import *

hsv_config = os.path.join(os.path.dirname(__file__), "..", "model", "thresholds.json")
try:
    with open(hsv_config, "r") as f:
        thresholds = json.load(f)
except FileNotFoundError:
    print(f"Error: Config file not found at {hsv_config}")
    thresholds = {}
    
hsv_config = build_hsv_map(thresholds)  

def position_color_center(image, color_name, min_area=1000, max_area=30000, model=0, iterations=1):
    """
    在图像中查找指定颜色的中心位置
    
    :param image: 输入图像
    :param color_name: 颜色名称，必须在配置文件中定义
    :param min_area: 最小轮廓面积 (仅用于 model 1)
    :param max_area: 最大轮廓面积 (仅用于 model 1)
    :param model: 使用的算法模型。0 或 1: 轮廓质心法, 2: 霍夫变换法
    :param iterations: 形态学操作迭代次数
    :return: 颜色中心坐标 (cx, cy) 或 None
    """
    try:
        if color_name not in hsv_config:
            raise ValueError(f"Color '{color_name}' not found in configuration.")
        # 将图像转换为HSV空间
        image = cv2.resize(image, (640, 480))  # 确保图像大小一致
        hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        lower = np.array(hsv_config[color_name]["lower"], dtype=np.uint8)
        upper = np.array(hsv_config[color_name]["upper"], dtype=np.uint8)
    except ValueError as e:
        print(f"Error: {e}")
        return None
    
    # 1. 颜色分割和形态学处理 (通用步骤)
    mask = cv2.inRange(hsv_image, lower, upper)
    # 使用闭运算填充线条内部的空洞
    kernel = np.ones((3,3), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=iterations)
    cv2.imwrite("/home/sunrise/cricket-training-questions/model/mask.jpg", mask)  # 显示掩码图像，便于调试

    # --- 模型 1: 轮廓质心法 (推荐) ---
    if model == 0 or model == 1:
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cv2.drawContours(image, contours, -1, (0, 255, 0), 2)  # 绘制轮廓
        if not contours:
            return None
        
        # 寻找面积最大的轮廓
        main_contour = max(contours, key=cv2.contourArea)
        area = cv2.contourArea(main_contour)

        # 面积过滤           
        if area < min_area or area > max_area:
            return None

        # 计算质心
        moments = cv2.moments(main_contour)
        if moments["m00"] == 0:
            return None
            
        cx = int(moments["m10"] / moments["m00"])
        cy = int(moments["m01"] / moments["m00"])
        # 返回质心坐标
        return (cx, cy)
    
    elif model == 2:

        canny_mask = cv2.Canny(mask, 50, 150, apertureSize=3)
        
        lines = cv2.HoughLinesP(canny_mask, 1, np.pi / 180, threshold=40, minLineLength=50, maxLineGap=20)
        if lines is None:
            print("No lines detected using Hough Transform.")
            return None
        # print(f"Detected {len(lines)} lines using Hough Transform.")
        for line in lines:
            x1, y1, x2, y2 = line[0]
            cv2.line(image, (x1, y1), (x2, y2), (0, 0, 255), 2)
      
        
        if lines is None:
            return None  

        horizontal_lines = []
        vertical_lines = []

        for line in lines:
            x1, y1, x2, y2 = line[0]
            angle = np.arctan2(y2 - y1, x2 - x1) * 180. / np.pi
            if abs(angle) < 45 or abs(angle) > 135:
                horizontal_lines.append(line)
            elif abs(abs(angle) - 90) < 45:
                vertical_lines.append(line)

        if not horizontal_lines or not vertical_lines:
            return None
        
        h_total_len = 0
        h_weighted_y_sum = 0

        for line in horizontal_lines:
            x1, y1, x2, y2 = line[0]
            length = np.sqrt((x2 - x1)**2 + (y2 - y1)**2)
            mid_y = (y1 + y2) / 2
            h_weighted_y_sum += mid_y * length
            h_total_len += length
    
        avg_y = h_weighted_y_sum / h_total_len if h_total_len > 0 else 0

        v_total_len = 0
        v_weighted_x_sum = 0

        for line in vertical_lines:
            x1, y1, x2, y2 = line[0]
            length = np.sqrt((x2 - x1)**2 + (y2 - y1)**2)
            mid_x = (x1 + x2) / 2
            v_weighted_x_sum += mid_x * length
            v_total_len += length

        avg_x = v_weighted_x_sum / v_total_len if v_total_len > 0 else 0

        if avg_x > 0 and avg_y > 0:
            return (int(avg_x), int(avg_y))
            
    return None



# # # --- 示例用法 ---
# if __name__ == '__main__':
#     # # 创建一个模拟的十字激光图像用于测试
#     # test_image = np.zeros((480, 640, 3), dtype=np.uint8)
#     # # 绘制一个绿色的十字
#     # cv2.line(test_image, (100, 240), (540, 240), (0, 255, 0), 4) # 水平线
#     # cv2.line(test_image, (320, 80), (320, 400), (0, 255, 0), 4)  # 垂直线
#     # # 期望中心点: (320, 240)
#     test_image = cv2.imread(r"C:\Users\19098\Desktop\OIP.webp")  # 替换为实际的测试图像路径
#     if test_image is None:
#         print("Error: Could not read the test image.")
#         exit(1)

#     # 使用模型1 (轮廓法)
#     center_model1 = position_color_center(test_image, 'green', model=1)
#     print(f"模型 1 (轮廓法) 检测到的中心点: {center_model1}")

#     # 使用模型2 (霍夫变换法)
#     center_model2 = position_color_center(test_image, 'green', model=2)
#     print(f"模型 2 (霍夫变换法) 检测到的中心点: {center_model2}")

#     if center_model1:
#         cv2.circle(test_image, center_model1, 5, (0, 0, 255), -1)
#         cv2.putText(test_image, f"Center: {center_model1}", (center_model1[0] + 15, center_model1[1]), 
#                     cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
        

#     # 在图像上绘制检测到的中心点以供可视化
#     if center_model2:
#         cv2.circle(test_image, center_model2, 10, (0, 0, 255), 2)
#         cv2.putText(test_image, f"Center: {center_model2}", (center_model2[0] + 15, center_model2[1]), 
#                     cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

#     cv2.imshow("Test Image with Detected Center", test_image)
#     cv2.waitKey(0)
#     cv2.destroyAllWindows()
